from .relys import BitSet
from .relys import Chip

class ALU(Chip):
    def __init__(self, width):
        super().__init__(width*3+5, [0, width*2+4], [width*2+4, width*3+5])
        self.TYP = self.pinIN@[0,3]
        self.EO = self.pinIN@[3]
        self.DAT1 = self.pinIN@[4,width+4]
        self.DAT2 = self.pinIN@[width+4, width*2+4]
        self.CF = self.pinOUT@[0]
        self._dat = BitSet(width+1)

@ALU.keyframe(1)
def _frame(chip):
    if chip.EO() : chip.pinOUT(chip._dat())
    else : chip.pinOUT(0)
    return chip

@ALU.keyframe(-1)
def _frame(chip):
    if   chip.TYP() == 0 : res = chip.DAT1() +  chip.DAT2()
    elif chip.TYP() == 1 : res = chip.DAT1() -  chip.DAT2()
    elif chip.TYP() == 2 : res = chip.DAT1() << chip.DAT2()
    elif chip.TYP() == 3 : res = chip.DAT1() >> chip.DAT2()
    elif chip.TYP() == 4 : res = chip.DAT1() &  chip.DAT2()
    elif chip.TYP() == 5 : res = chip.DAT1() |  chip.DAT2()
    elif chip.TYP() == 6 : res = chip.DAT1() ^  chip.DAT2()
    elif chip.TYP() == 7 : res = ~chip.DAT1()
    chip._dat(res)
    if res > BitSet(chip._dat.width-1)(-1) : chip.CF(1)
    else : chip.CF(0)
    return chip

# module test
if __name__ == "__main__":
    alu = ALU(8)
    alu.TYP(0)
    alu.DAT1(0b01010101)
    alu.DAT2(0b10101011)
    alu.EO(1)
    for i in range(ALU.maxframe+1) : 
        alu.update(i)
        print(f"<frame>:[{i}]", alu)
